import pandas as pd
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']
import os
import tushare as ts

ts.set_token('084256242bd0568c9ea0786ff7019ea73e1d83efa7d51a3454f99d67')
pro = ts.pro_api()


def get_data(code, start, end):
    df = pro.daily(ts_code=code, autype='qfq', start_date=start, end_date=end)
    print(df)
    df.index = pd.to_datetime(df.trade_date)
    # 设置把日期作为索引
    # df['ma'] = 0.0  # Backtrader需要用到
    # df['openinterest'] = 0.0  # Backtrader需要用到
    # 定义两个新的列ma和openinterest
    df = df[['open', 'high', 'low', 'close', 'vol']]
    # 重新设置df取值，并返回df
    return df


def acquire_code(inp_code):  # 只下载一只股票数据，且只用CSV保存   未来可以有自己的数据库
    inp_start = input("请输入该股票数据集获取的开始时间:'\n'")
    inp_end = input("请输入该股票数据集获取的结束时间:'\n'")
    df = get_data(inp_code, inp_start, inp_end)
    print(df.info())
    # 输出统计各列的数据量
    print("—" * 30)
    # 分割线
    print(df.describe())
    # 输出常用统计参数
    df.sort_index(inplace=True)
    # 把股票数据按照时间正序排列
    path = os.path.join(os.path.join(os.getcwd(),
                                     "D:\deep learning project\lstm-stock_prediction"), inp_code + ".csv")
    # os.path地址拼接，''数据地址''为文件保存路径
    # path = os.path.join(os.path.join(os.getcwd(),"数据地址"),inp_code+"_30M.csv")
    df.to_csv(path)


stock_code = input("请输入想要预测的股票代码:\n")
acquire_code(stock_code)

import csv

# 打开CSV文件
with open(stock_code+'.csv', 'r') as file:
    # 读取CSV文件的内容
    csv_data = csv.reader(file)

    # 将读取的数据转换为列表
    data_list = list(csv_data)

    # 确定要修改的行数（索引从0开始）
    row_index = 0

    # 修改指定行的数据
    data_list[row_index] = ['Data', 'Open', 'High', 'Low', 'Close', 'Vol']

    # 将修改后的数据写回CSV文件
    with open(stock_code+'.csv', 'w', newline='') as file:
        writer = csv.writer(file)
        writer.writerows(data_list)
